import pandas as pd
import os

def calculate_macd_for_stock(df):
    # 添加新列
    df['ema12'] = 0.0
    df['ema26'] = 0.0
    df['DIF'] = 0.0
    df['DEA'] = 0.0
    df['MACD'] = 0.0

    # 初始化第一条记录
    df.at[0, 'ema12'] = df.at[0, '周线_后复权收盘价']
    df.at[0, 'ema26'] = df.at[0, '周线_后复权收盘价']
    df.at[0, 'DIF'] = 0.0
    df.at[0, 'DEA'] = 0.0
    df.at[0, 'MACD'] = 0.0

    # 从第二条记录开始计算
    for i in range(1, len(df)):
        prev_ema12 = df.at[i-1, 'ema12']
        prev_ema26 = df.at[i-1, 'ema26']
        prev_dea = df.at[i-1, 'DEA']
        current_close = df.at[i, '周线_后复权收盘价']

        # 计算 EMA12 和 EMA26
        df.at[i, 'ema12'] = (2/13) * current_close + (11/13) * prev_ema12
        df.at[i, 'ema26'] = (2/27) * current_close + (25/27) * prev_ema26

        # 计算 DIF
        df.at[i, 'DIF'] = df.at[i, 'ema12'] - df.at[i, 'ema26']

        # 计算 DEA
        df.at[i, 'DEA'] = (2/10) * df.at[i, 'DIF'] + (8/10) * prev_dea

        # 计算 MACD
        df.at[i, 'MACD'] = 2 * (df.at[i, 'DIF'] - df.at[i, 'DEA'])

    return df

def main():
    path = r'D:\pythonProject\BaiduSyncdisk\LongMoneyTrade\各部分数据\后复权周线'

    # 列出所有CSV文件
    csv_files = [f for f in os.listdir(path) if f.endswith('.csv')]
    print(f"发现以下CSV文件：{csv_files}")

    for filename in csv_files:
        file_path = os.path.join(path, filename)
        print(f"正在处理文件：{filename}...")

        try:
            df = pd.read_csv(file_path)

            if df.empty:
                print(f"{filename} 为空，跳过。")
                continue

            if '周线_后复权收盘价' not in df.columns:
                print(f"{filename} 缺少必要字段 '周线_后复权收盘价'，跳过。")
                continue

            updated_df = calculate_macd_for_stock(df)

            # 保留两位小数
            columns_to_round = ['ema12', 'ema26', 'DIF', 'DEA', 'MACD']
            updated_df[columns_to_round] = updated_df[columns_to_round].round(2)

            # 写回原文件
            updated_df.to_csv(file_path, index=False)
            print(f"已完成处理：{filename}")

        except Exception as e:
            print(f"处理 {filename} 时发生错误：{e}")
            import traceback
            traceback.print_exc()

if __name__ == '__main__':
    main()
